package com.ruoyi.web.handler;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.*;
import java.util.concurrent.CopyOnWriteArrayList;
@Component
public class SpringWebSocketHandler extends TextWebSocketHandler {
    private final CopyOnWriteArrayList<WebSocketSession> sessions = new CopyOnWriteArrayList<>();
    private final ByteArrayOutputStream baos = new ByteArrayOutputStream();
    private final PrintStream originalOut = System.out;
    private final PrintStream printStream = new PrintStream(new OutputStream() {
        private final OutputStream systemOut = originalOut;
        private final OutputStream byteArrayOut = baos;

        @Override
        public void write(int b) {
            try {
                systemOut.write(b);
                byteArrayOut.write(b);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
        @Override
        public void flush() {
            try {
                systemOut.flush();
                byteArrayOut.flush();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }

        @Override
        public void close() {
            try {
                systemOut.close();
                byteArrayOut.close();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }) {
        @Override
        public void println(String x) {
            super.println(x);
            broadcast(x + "\n");
        }
    };
    public SpringWebSocketHandler() {
        // Redirect System.out to our custom PrintStream
        System.setOut(printStream);
    }
    /**
     * js调用websocket.send时候，会调用该方法
     * */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        if ("start".equals(message.getPayload())) {
            sessions.add(session);
        } else if ("stop".equals(message.getPayload())) {
            sessions.remove(session);
        }
    }
    /**
     * 广播消息给所有连接的客户端
     */
    private void broadcast(String message) {
        for (WebSocketSession session : sessions) {
            try {
                if (session.isOpen()) { // 检查会话是否打开
                    synchronized (session) { // 加锁确保线程安全
                        session.sendMessage(new TextMessage(message));
                    }
                }
            } catch (IOException e) {
                // 处理异常，例如移除已关闭的会话
                sessions.remove(session);
                System.err.println("Failed to send message to session " + session.getId() + ": " + e.getMessage());
            }
        }
    }
    /**
     * 允许外部调用发送消息给所有连接的客户端
     */
    public void sendMessageToAll(String message) {
        broadcast(message);
    }
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 当连接建立时调用
        super.afterConnectionEstablished(session);
        sessions.add(session); // 在连接建立时也添加session
    }
    /**
     * 关闭连接时触发
     * */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        sessions.remove(session);
        super.afterConnectionClosed(session, status);
    }
}
